package karpiel.boguchwal.wojciech.przemek

import karpiel.boguchwal.wojciech.przemek.Ast._

object TypeInferer {
  def inferType(expr: Ast)(implicit ctx: Context, cache: Cache): Ast = inferTypeNoNorm(expr.normalized)(ctx, cache)

  private[przemek] def inferTypeNoNorm(expr: Ast)(implicit ctx: Context, cache: Cache): Ast = {
    cache.infer(expr,
      () => expr match {
      case UnitValue => UnitType
      case UnitType => Universe(0)
      case BoolType => Universe(0)
      case BoolTrue => BoolType
      case BoolFalse => BoolType
      case BoolElim(ifTrue, ifFalse) =>
        ifTrue.inferCC.assertEqualsOpt(ifFalse.inferCC)
        PiType(BoolType, ifTrue.inferCC)
      case BoolRec(boolToType, ifTrue, ifFalse) =>
        val pi = boolToType.getAs[PiValue]
        pi.typeT.assertEqualsOpt(BoolType)
        pi.value.inferType(ctx.append(BoolType), implicitly).getAs[Universe](ctx.append(BoolType), implicitly)
        ifTrue.inferCC.assertEqualsOpt(PiElim(boolToType, BoolTrue))
        ifFalse.inferCC.assertEqualsOpt(PiElim(boolToType, BoolFalse))
        PiType(BoolType, pi.value)
      case Universe(i) => Universe(i + 1)
      case b: Bruijn => ctx.getType(b)
      case PiType(typeT, value) =>
        val u1 = typeT.inferCC.getAs[Universe]
        val u2 = value.inferType(ctx.append(typeT), implicitly).getAs[Universe]
        Universe(Math.max(u1.i, u2.i))
      case PiValue(typeT, value) =>
        typeT.inferCC.getAs[Universe]
        PiType(typeT, value.inferType(ctx.append(typeT), implicitly))
      case PiElim(piVal, argument) =>
        piVal.normalized match {
          case piValue: PiValue =>
            val argumentType = argument.inferCC
            if (piValue.typeT != argumentType) throw new TypecheckException()
            val functionBodyContext = ctx.append(argumentType, argument)
            val piValueTypeInExtCtx = piValue.value.inferType(functionBodyContext, implicitly).normalized(functionBodyContext, implicitly)
            if (Utils.hasBruijn(piValueTypeInExtCtx, 0)) {
              PiElim(piValue, argument)
            } else {
              Utils.unshift(piValueTypeInExtCtx, delta = 1)
            }
          case b: BoolElim => //todo skopiowany kod
            argument.inferCC.assertEqualsOpt(BoolType)
            b.ifFalse.inferCC.assertEqualsOpt(b.ifTrue.inferCC)
            b.ifTrue.inferCC
          case e@EmptyElim(anyType) =>
            e.inferCC
            argument.inferCC.assertEqualsOpt(EmptyType)
            anyType
          case piValUnnorm =>
            val tmp_ = piValUnnorm.inferCC
            val piType = tmp_.getAs[PiType]
            //todo tu nie można po prostu dać piType.value
            //bo to żyje w innym kontelście
            // anie można przesunąć kontekstu bo może być bruijn0
            val argType = argument.inferCC
            piType.typeT.assertEqualsOpt(argType) //todo w innych miejscach też?
          val piValNormExtendedCtx = piType.value.normalized(ctx.append(argType, argument), implicitly)
            if (Utils.hasBruijn(piValNormExtendedCtx, 0)) {
              PiElim(piValUnnorm, argument.normalized)
            } else {
              Utils.unshift(piValNormExtendedCtx, delta = 1)
            }
        }
      case SumType(aType, bType) =>
        val uA = aType.inferCC.getAs[Universe]
        val uB = bType.inferCC.getAs[Universe]
        Universe(Math.max(uA.i, uB.i))
      case SumInl(aValue, bType) => SumType(aValue.inferCC, bType)
      case SumInr(aType, bValue) => SumType(aType, bValue.inferCC)
      case SumElim(aFn, bFn) =>
        val aPi = aFn.getAs[PiValue]
        val bPi = bFn.getAs[PiValue]
        val inType = SumType(aPi.typeT, bPi.typeT)
        val aPiType = aPi.inferCC.asInstanceOf[PiType]
        val bPiType = bPi.inferCC.asInstanceOf[PiType]
        val newCtx = ctx.append(inType)
        aPiType.value.assertEquals(bPiType.value)(newCtx, implicitly)
        PiType(inType, aPiType.value.normalized(newCtx, implicitly))
      case SigmaType(pi: Ast) =>
        val piVal = pi.getAs[PiValue]
        val u1 = piVal.typeT.inferCC.getAs[Universe]
        val newCtx = ctx.append(piVal.typeT.normalized)
        val u2 = piVal.value.normalized(newCtx, implicitly).inferType(newCtx, implicitly).getAs[Universe]
        Universe(Math.max(u1.i, u2.i))
      case SigmaValue(x, aX, aType) =>
        val xType = x.inferCC.normalized
        PiElim(aType, x).assertEquals(aX.inferCC)
        val pt = aType.inferType.getAs[PiType]
        pt.typeT.assertEquals(xType)
        SigmaType(aType)
      case SigmaElim(pi) =>
        pi.normalized match {
          case PiValue(xType, PiValue(aXType, ex)) =>
            PiValue(SigmaType(PiValue(xType, aXType)), Utils.sigmaElimHelper(ex, 0)).inferCC
          case any => throw new TypecheckException()
        }
      case Proj1(sigma) =>
        sigma.normalized match {
          case sigmaV: SigmaValue =>
            sigmaV.inferCC //safety check
            sigmaV.x.inferCC
          case anyElse =>
            anyElse.inferCC.getAs[SigmaType].pi.getAs[PiValue].typeT
        }
      case Proj2(sigma) =>
        sigma.normalized match {
          case sigmaV: SigmaValue =>
            sigmaV.inferCC //safety check
            sigmaV.aX.inferCC
          case anyElse =>
            val a = anyElse.inferCC.getAs[SigmaType]
            val b = a.pi.getAs[PiValue]
            val inferType = Proj1(anyElse).inferCC
            inferType.assertEqualsOpt(b.typeT)
            Utils.unshift(b.value.normalized(ctx.append(inferType, Proj1(anyElse)), implicitly), 0, 1)
        }
      case IdType(x, y) =>
        x.inferCC.assertEqualsOpt(y.inferCC)
        val u1 = x.inferCC.inferCC.getAs[Universe]
        val u2 = y.inferCC.inferCC.getAs[Universe]
        Universe(Math.max(u1.i, u2.i))
      case Refl(x) => IdType(x, x)
      case J(typeFunc, reflFunc) =>
        typeFunc.normalized match {
          case PiValue(a1T, PiValue(a2T, PiValue(idXYT, type_))) =>
            val a1TN = a1T.normalized
            val a2TN = a2T.normalized(ctx.append(a1TN), implicitly)
            if (Utils.hasBruijn(a2TN, 0)) {
              throw new TypecheckException("Што делать?")
            } else {
              if (Utils.unshift(a2TN, delta = 1) != a1T.normalized) {
                throw new TypecheckException()
              }
            }
            val idCtx = ctx.append(a1TN).append(a2TN)
            val id = idXYT.getAs[IdType](idCtx, implicitly)
            id.x.assertEquals(Bruijn(1))(idCtx, implicitly) //todo normalizacja to moze zepsuć
            id.y.assertEquals(Bruijn(0))(idCtx, implicitly)
            val typeN = type_.normalized(idCtx.append(id), implicitly)

            reflFunc.inferCC match { //only check
              case PiType(typeT, valueType) =>
                typeT.assertEqualsOpt(a1T)
                //znowu porównywanie między kontekstami :(
                val typZWartościP1 = valueType.normalized(ctx.append(a1TN), implicitly)

                val typPodantP3 = type_.normalized(
                  ctx.append(a1TN).append(Utils.shift(a1TN, delta = 1), Bruijn(0))
                    .append(Ast.IdType(Bruijn(1), Bruijn(0)), Refl(Bruijn(1))), implicitly)
                if (Utils.hasBruijn(typPodantP3, 0) || Utils.hasBruijn(typPodantP3, 1)) {
                  throw new TypecheckException("ZAL")
                }
                val unshifted = Utils.unshift(typPodantP3, delta = 2)
                if (unshifted != typZWartościP1) {
                  throw new TypecheckException("NIERÓWNASIE:\n"+unshifted+"\n"+typZWartościP1)
                }

              case _ => throw new TypecheckException("tak naprawdę to nie")
            }
            PiType(a1TN, PiType(a2TN, PiType(IdType(Bruijn(1), Bruijn(0)), typeN)))
          case _ =>
            throw new TypecheckException("tak naprawdę to nie")
        }
      case EmptyType => Universe(0)
      case EmptyElim(anyType) =>
        anyType.inferCC.getAs[Universe]
        PiType(EmptyType, Utils.shift(anyType, 0, 1))
      case WType(arityFn) =>
        val piType = arityFn.inferCC.getAs[PiType]
        val u1 = piType.typeT.inferCC.getAs[Universe]
        val ctxExt = ctx.append(piType.typeT)
        val u2 = piType.value.getAs[Universe](ctxExt, implicitly)
        Universe(Math.max(u1.i, u2.i))
      case WSup(wType_, label, fn) =>
        val wType = wType_.getAs[WType]
        val wPiValue = wType.arityFn.getAs[PiValue]
        label.inferCC.assertEqualsOpt(wPiValue.typeT)
        val BaType = PiElim(wPiValue, label).normalized
        BaType.inferCC.getAs[Universe]
        val fnPi = fn.inferCC.getAs[PiType]
        fnPi.typeT.assertEqualsOpt(BaType)
        if (fnPi.value.normalized(ctx.append(BaType), implicitly) != wType)
          throw new TypecheckException()
        wType

      case WElim(e, eSup) => //todo this would probably fail on legit cases, this is too strict
        //check that E is okay
        val b = e.getAs[PiValue]
        val bType = b.inferCC.getAs[PiType]
        bType.value.getAs[Universe](ctx.append(b.typeT), implicitly)
        val wType = b.typeT.getAs[WType]
        wType.inferCC
        val wPi = wType.arityFn.getAs[PiValue]
        //check Esup
        /*
         Pi(a:A,
         Pi (u: B(a) -> WxBx,
    *    Pi(q: Pi(y:B(a), E(u(y))),
    *    E(sup(a,u)))))))
         */
        val expectedType: Ast =
        PiType(
          wPi.typeT //a:A
          ,
          PiType(
            PiType(PiElim(wType.arityFn.shiftn(1), Bruijn(0)), wType.shiftn(2)) //u :B(a) -> WxBx
            ,
            PiType(
              PiType(PiElim(wType.arityFn.shiftn(2), Bruijn(1)), PiElim(e, PiElim(Bruijn(1), Bruijn(0)))) //Pi(y: B(a), E(u(y)))
              ,
              PiElim(e, WSup(wType.shiftn(2), Bruijn(2), Bruijn(1))) //E(sup(a,u)
            )
          )
        )
        val eSupType = eSup.inferCC
        val eSupTypen = eSupType.normalized
        val expectedTypeN = expectedType.normalized
        eSupTypen.assertEqualsOpt(expectedTypeN)
        //result
        PiType(wType, PiElim(Utils.shift(e, 0, 1), Bruijn(0))).normalized
      case _: NamedType => throw new TypecheckException()
      case _: Var => throw new TypecheckException()
    })
  }
}
